#%%
import matplotlib.pyplot as plt
import numpy as np, tqdm, pandas as pd, sys, os, pathlib
import scipy.odr as odr
wd = str(pathlib.Path(os.path.dirname(__file__)).parent)
print(wd)
sys.path.append(wd)
import raw_exp_data

# Data input
idir = './image_slope_measurements_data/'

###########################
# Perspective Calibration #
###########################

runs = [
    'e6',
    'e7',
    'e8',
    'e11',
    'e12',
        ]

persp_data = dict([])

for run in runs:
    data = dict([])
    test = pd.read_excel(idir+'gb_s3_'+run+'_persp.xlsx')
    data['pix'] = test[test.Unit=='pixels'].Length.values[0]
    xs = []
    ls = []
    for i in range(4):
        xs.append(test[test.Measurement=='x_%s'%(i+1)].Length.values[0])
        ls.append(test[test.Measurement=='%s'%(i+1)].Length.values[0])
    data['x'] = np.array(xs)
    data['l'] = np.array(ls)/ls[1]    # The 'truth' is the second one
    persp_data[run]= data
    
aa = []
bb = []
cc  = []
for ii,run in enumerate(runs):
    x = persp_data[run]['x']
    l = persp_data[run]['l']
    plt.plot(x,l,'o',label=run,color=(0,0,(ii+1)/len(runs),1))
    ## Fitting
    # Linear fit
    def func(B, x): return B[0]*(x-B[1])**(2)+B[2]
    pwrlaw = odr.Model(func)
    mydata = odr.RealData(x,l)
    myodr=odr.ODR(mydata,pwrlaw,beta0=[0.0,0.0,0.0])
    myoutput=myodr.run()
    a = myoutput.beta[0]
    da =myoutput.sd_beta[0]
    xshift =  myoutput.beta[1]
    dxshift = myoutput.sd_beta[1]
    c =  myoutput.beta[2]
    dc = myoutput.sd_beta[2]

    aa.append(a)
    bb.append(xshift)
    cc.append(c)
    
    print(run,a,'xshift = ',xshift,c)

    xpl = np.linspace(np.min(x),np.max(x))
    plt.plot(xpl,func([a,xshift,c],xpl),'--',color=(0,0,(ii+1)/len(runs),1))

# Taking average of e7-e12
am = np.mean(aa[1:])
bm = np.mean(bb[1:])
cm = np.mean(cc[1:])
print(am,bm,cm)
plt.plot(xpl,am*(xpl-50)**2+cm,'-r',lw=2)

for ii,run in enumerate(runs):
    x = persp_data[run]['x']
    l = persp_data[run]['l']
    l = l[1]*l/(am*(x-50)**2+cm)
    plt.plot(x,l,'-o',label=run+' corrected',color=(0,(ii+1)/len(runs),0,1))
plt.legend(loc=(1.01,0.25))
plt.title("Calibration")
plt.show()

# Calculating the xshift for gb_s3_e{1-4}
xshift7 = bm
xshift6 = 254.257952
dx_plug7 = 540.
dx_plug6 = 621.2

m,b = np.polyfit([dx_plug7,dx_plug6],[xshift7,xshift6],1)
print(m,b)

dx_plug3 = 510.14

xshift3 = m*dx_plug3+b
print('xshift3 = %f' % xshift3)

# Quality check:
for ii,run in enumerate(['e3','e6','e7']):
    imdat = pd.read_excel(idir+'gb_s3_'+run+'_v2.xlsx')
    len_list = len(imdat[imdat.Measurement.str.contains('_b')].Length.values)
    bed_h_im = []
    water_h_im = []
    cal_h_im = []
    dxs = [0.0]
    for i in range(len_list):
        bed_h_im.append(imdat[imdat.Measurement=='%s_b' % (i+1)].Length.values[0])
        water_h_im.append(imdat[imdat.Measurement=='%s_w' % (i+1)].Length.values[0])
        cal_h_im.append(imdat[imdat.Measurement=='%s_c' % (i+1)].Length.values[0])
        if i <(len_list-1):
            dxs.append(imdat[imdat.Measurement=='%s%s_dx' % (i+1,i+2)].Length.values[0])
        else:
            pass

    # Include lens corrections
    x = np.array(dxs)
    bed_h_im = np.array(bed_h_im)
    water_h_im = np.array(water_h_im)
    l = np.array(cal_h_im)
    plt.plot(x,l,'-o',color=(0,0,(ii+1)/len(runs),1))
    if run in ['e6','e5']:
        xshift = xshift6
    elif run in ['e1','e2','e3','e4']:
        xshift = xshift3
    else:
        xshift = xshift7
    l = l/(func([am,xshift,cm],x))
    plt.plot(x,l,'-o',color=(0,(ii+1)/len(runs),0,1),label = run)
plt.legend(loc=(1.01,0.25))
plt.title("Quality check for calibration")
plt.show()

###

aa = []
bb = []
cc  = []
for ii,run in enumerate(runs):
    x = persp_data[run]['x']
    l = persp_data[run]['l']
    plt.plot(x,l,'o',label=run,color=(0,0,(ii+1)/len(runs),1))
    xpl = np.linspace(np.min(x),np.max(x))
    if run in ['e6','e5']:
        xshift = xshift6
    elif run in ['e1','e2','e3','e4']:
        xshift = xshift5
    else:
        xshift = xshift7
    
    plt.plot(xpl,func([am,xshift,cm],xpl),'--',color=(0,0,(ii+1)/len(runs),1))

plt.legend()

for ii,run in enumerate(runs):
    x = persp_data[run]['x']
    l = persp_data[run]['l']
    if run in ['e6','e5']:
        xshift = xshift6
    elif run in ['e1','e2','e3','e4']:
        xshift = xshift5
    else:
        xshift = xshift7
    l = l/(func([am,xshift,cm],x))
    plt.plot(x,l,'-o',label=run+' corrected',color=(0,(ii+1)/len(runs),0,1))
    
plt.plot(xpl,func([am,xshift3,cm],xpl),'-r')
plt.title("Quality check number 2")
plt.legend(loc=(1.01,0.25))
plt.show()

#######################
# Slope measurements  #
#######################

flume_slope = 35.4/1000-10/1000

# Data of all runs
full_data = dict([])

# Data of runs that don't have 'true' measurements with rulers
export_data = dict([])

# All the runs
runs = [
    'gb_s3_e1',
    'gb_s3_e2',
    'gb_s3_e3',
    'gb_s3_e4',
    'gb_s3_e5',
    'gb_s3_e6',
    'gb_s3_e7',
    'gb_s3_e8',
    'gb_s3_e9',
    'gb_s3_e10',
    'gb_s3_e11',
    'gb_s3_e12',
    'ng_s3_e2',
    'ng_s3_e3',
]

for run in runs:
    # Split run names:
    exp = run.split('_')[2]
    series = run.split('_')[0]+'_'+run.split('_')[1]
    
    data_run = dict([])
    
    # Add 'true' data taken from rulers
    slope_inds = np.array(raw_exp_data.series_data[series][exp]['slope_inds'])
    xs = np.array(raw_exp_data.series_data[series][exp]['meas_dist_along_flume'])[slope_inds[0]:slope_inds[1]]
    bed_h = np.array(raw_exp_data.series_data[series][exp]['bed_h_range_equil'])[slope_inds[0]:slope_inds[1]]
    water_h = np.array(raw_exp_data.series_data[series][exp]['water_h_range_equil'])[slope_inds[0]:slope_inds[1]]
    
    try: # If nan is encountered, this will give an error
        m,b = np.polyfit(xs,bed_h,1)
        data_run['bed_real'] = m+flume_slope

        m,b = np.polyfit(xs,water_h,1)    
        data_run['water_real'] = m+flume_slope
    except:
        data_run['bed_real'] = np.nan
 
        data_run['water_real'] = np.nan
        
    
    # Now measure from images
    imdat = pd.read_excel(idir+run+'_v2.xlsx')
    len_list = len(imdat[imdat.Measurement.str.contains('_b')].Length.values)
    bed_h_im = []
    water_h_im = []
    dxs = [0.0]
    for i in range(len_list):
        bed_h_im.append(imdat[imdat.Measurement=='%s_b' % (i+1)].Length.values[0])
        water_h_im.append(imdat[imdat.Measurement=='%s_w' % (i+1)].Length.values[0])
        if i <(len_list-1):
            dxs.append(imdat[imdat.Measurement=='%s%s_dx' % (i+1,i+2)].Length.values[0])
        else:
            pass

    # Include lens corrections
    if run in ['gb_s3_e6','gb_s3_e5']:
        xshift = xshift6
    elif run in ['gb_s3_e1','gb_s3_e2','gb_s3_e3','gb_s3_e4']:
        xshift = xshift3
    else:
        xshift = xshift7
        
    xs_im = np.array(dxs)
    # Making sure to add correction due to perspective using the 'func' function
    bed_h_im = np.array(bed_h_im)/(func([am,xshift,cm],xs_im))
    water_h_im = np.array(water_h_im)/(func([am,xshift,cm],xs_im))
    
    m,b = np.polyfit(xs_im,bed_h_im,1)
    data_run['bed_im'] = m+flume_slope
    
    m,b = np.polyfit(xs_im,water_h_im,1)
    data_run['water_im'] = m+flume_slope
    
    # Calculate percent error:
    pce = abs((data_run['bed_real']-data_run['bed_im'])/data_run['bed_real'])*100
    data_run['pe_bed'] = pce
    
    
    pce = abs((data_run['water_real']-data_run['water_im'])/data_run['water_real'])*100
    data_run['pe_water'] = pce
    
    # Save run to the full data:
    full_data[run] = data_run
    
    # Only put those that don't have ruler measurements into data that will be exported
    print(bed_h, water_h)
    if len(bed_h) == 0 or len(water_h) == 0:
        export_data[run] = data_run

################################
# Comparison, Error estimation #
################################

# First, error estimation, using percent error.

types = [
    'bed',
#     'water',
]
mtype = ['o','x']

plt.figure(figsize=(8,6))

xdat = []
ydat = []
for ii,ttype in enumerate(types):
    for run in full_data:
        exp = run.split('_')[2]
        series = run.split('_')[0]+'_'+run.split('_')[1]
        if run=='gb_s3_e6':
            if exp!='e7' and exp!='e9' and series!='ng_s3':
                xdat.append(full_data[run][ttype+'_real'])
                ydat.append(full_data[run]['pe_'+ttype])
            plt.scatter(full_data[run][ttype+'_real'],full_data[run]['pe_'+ttype],color='k',marker=mtype[ii],label=ttype)
#             if ttype=='water':
            plt.annotate(run,(full_data[run][ttype+'_real']+0.001,full_data[run]['pe_'+ttype]))
        else:
            if exp!='e7' and exp!='e9' and series!='ng_s3':
                xdat.append(full_data[run][ttype+'_real'])
                ydat.append(full_data[run]['pe_'+ttype])
            plt.scatter(full_data[run][ttype+'_real'],full_data[run]['pe_'+ttype],color='k',marker=mtype[ii])
#             if ttype=='water':
            plt.annotate(run,(full_data[run][ttype+'_real']+0.001,full_data[run]['pe_'+ttype]))

# Remove nans
xdat = np.array(xdat)
ydat = np.array(ydat)
xdat = xdat[~np.isnan(xdat)]
ydat = ydat[~np.isnan(ydat)]
m_pce,b_pce = np.polyfit(xdat,np.log10(ydat),1)
xpl = np.linspace(np.min(xdat),np.max(xdat),100)
plt.plot(xpl,10**(m_pce*xpl+b_pce),'--r',label='exponential fit')
plt.xlabel('Slope',fontsize=20)
plt.ylabel('Percent Error',fontsize=20)
plt.legend(fontsize=20)
plt.show()

# Now to look for a systematic difference:
types = [
    'bed',
    'water',
]
mtype = ['o','x']

plt.figure(figsize=(8,6))


mean_offset = dict([])
for ii,ttype in enumerate(types):
    xdat = []
    ydat = []
    for run in full_data:
        exp = run.split('_')[2]
        series = run.split('_')[0]+'_'+run.split('_')[1]
        if run=='gb_s3_e6':
            xdat.append(full_data[run][ttype+'_real'])
            ydat.append(full_data[run][ttype+'_real']-full_data[run][ttype+'_im'])
            plt.scatter(full_data[run][ttype+'_real'],full_data[run][ttype+'_real']-full_data[run][ttype+'_im'],color='k',marker=mtype[ii],label=ttype)
#             if ttype=='water':
#             plt.annotate(run,(full_data[run][ttype+'_real']+0.001,full_data[run][ttype+'_real']-full_data[run][ttype+'_im']))
        else:
            xdat.append(full_data[run][ttype+'_real'])
            ydat.append(full_data[run][ttype+'_real']-full_data[run][ttype+'_im'])
            plt.scatter(full_data[run][ttype+'_real'],full_data[run][ttype+'_real']-full_data[run][ttype+'_im'],color='k',marker=mtype[ii])
#             if ttype=='water':
#             plt.annotate(run,(full_data[run][ttype+'_real']+0.001,full_data[run][ttype+'_real']-full_data[run][ttype+'_im']))
    # Remove nans
    xdat = np.array(xdat)
    ydat = np.array(ydat)
    xdat = xdat[~np.isnan(xdat)]
    ydat = ydat[~np.isnan(ydat)]
    mean_offset[ttype] = np.mean(ydat)
    print("Mean offset, %s = %f" % (ttype,np.mean(ydat)))


plt.xlabel('Real Slope',fontsize=20)
plt.ylabel('Real-Image Estimate',fontsize=20)
plt.legend(fontsize=20)
plt.show()

###############################################
# Final export of data for image measurements #
###############################################

data_output = dict([])
for run in export_data:
    data1 = dict([])
    print('bed',mean_offset['bed'])
    print(run,export_data[run]['bed_im'],export_data[run]['bed_im'])
    data1['Sb'] = export_data[run]['bed_im']
    data1['dSb'] = data1['Sb']*2*10**(m_pce*data1['Sb']+b_pce)/100
    
    print('water',mean_offset['water'])
    print(run,export_data[run]['water_im'],export_data[run]['water_im'])
    data1['Sw'] = export_data[run]['water_im']
    data1['dSw'] = data1['Sw']*2*10**(m_pce*data1['Sw']+b_pce)/100
    
    if run=='gb_s3_e1':
        data1['d'] = np.mean(water_h_im[:-2]-bed_h_im[:-2])
        data1['dd'] = np.std(water_h_im[:-2]-bed_h_im[:-2])/np.sqrt(len(water_h_im[:-2]-bed_h_im[:-2]))
    else:
        data1['d'] = np.mean(water_h_im-bed_h_im)
        data1['dd'] = np.std(water_h_im-bed_h_im)/np.sqrt(len(water_h_im-bed_h_im))
    
    data_output[run]=data1
    
df = pd.DataFrame(data=data_output)

df.to_csv(wd+'/image_slope_meas.csv')
# %%
